import os

import numpy as np
from utils.metrics import R1_mAP, R1_mAP_Pseudo, eval_func
from datasets.mydata import mydata
if __name__ == "__main__":

    base_dir = '../json/'
    model1 = 'efficientnet-b2'  # 92 95
    model2 = 'resnet50_ibn_a'  # 93 96
    model3 = 'regnet/normal'  # 92 95
    model4 = 'hrnet'  # 87 93
    model5 = 'resnet50_ibn_a/pseudo'  # 92 953
    model6 = 'regnet/pseudo'  # 92 953

    dataset = mydata('../data')
    query_data = dataset.query_normal
    gallery_data = dataset.gallery_normal
    query, q_pids, q_camids = [],[],[]
    gallery, g_pids, g_camids = [], [],[]
    for (path,pid,camid) in query_data:
        query.append(path)
        q_pids.append(pid)
        q_camids.append(camid)
    for (path, pid, camid) in gallery_data:
        gallery.append(path)
        g_pids.append(pid)
        g_camids.append(camid)
    q_pids,q_camids,g_pids,g_camids = np.asarray(q_pids),np.asarray(q_camids),np.asarray(g_pids),np.asarray(g_camids)

    # distmat = np.load(base_dir + model1 + '/distmat_v.npy')  # 各个模型的距离矩阵相加
    # distmat += np.load(base_dir + model2 + '/distmat_v.npy')
    # distmat += np.load(base_dir + model3 + '/distmat_v.npy')
    # distmat = np.load(base_dir + model4 + '/distmat_v.npy')
    # distmat = np.load(base_dir + model5 + '/distmat_v.npy')
    distmat = np.load(base_dir + model3 + '/distmat_v.npy')

    cmc, mAP = eval_func(distmat, q_pids, g_pids, q_camids, g_camids, 200)
    print("mAP:{:.3f}\nrank1:{:.3f}\nrank5:{:.3f}".format(mAP, cmc[0], cmc[5]))